package com.abc.utils;

import com.abc.config.SparkGptConfig;
import com.abc.constant.SparkConstant;
import com.abc.dto.request.ApiRequestRequestDto;
import com.abc.dto.response.SseResponse;
import okhttp3.*;
import com.alibaba.fastjson.JSON;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;

import java.io.BufferedReader;
import java.io.IOException;
import java.io.InputStreamReader;

@Component
public class SparkRequestWebSocketUtil {
    private static final Logger logger = LoggerFactory.getLogger(SparkRequestWebSocketUtil.class);
    private final SparkGptConfig sparkGptConfig;
    private final OkHttpClient client = new OkHttpClient();

    public SparkRequestWebSocketUtil(SparkGptConfig sparkGptConfig) {
        this.sparkGptConfig = sparkGptConfig;
    }

    public void sendApiStreamRequest(ApiRequestRequestDto requestDto, WebSocketSession session) {
        MediaType mediaType = MediaType.parse("application/json; charset=utf-8");
        RequestBody requestBody = RequestBody.create(JSON.toJSONString(requestDto), mediaType);

        Request request = new Request.Builder()
                .url(SparkConstant.CHAT_URL)
                .addHeader("Authorization", SparkConstant.AUTHORIZATION + sparkGptConfig.getApiPassword())
                .post(requestBody)
                .build();

        client.newCall(request).enqueue(new Callback() {
            @Override
            public void onFailure(Call call, IOException e) {
                logger.error("请求因网络错误失败", e);
                // 处理错误，例如发送错误信息到客户端
            }

            @Override
            public void onResponse(Call call, Response response) {
                if (response.isSuccessful()) {
                    try (ResponseBody responseBody = response.body()) {
                        if (responseBody != null) {
                            handleWebSocketResponse(responseBody, session);
                        }
                    } catch (IOException e) {
                        logger.error("读取响应时出错", e);
                    }
                } else {
                    logger.warn("请求失败，状态码: {}", response.code());
                    // 处理请求失败的逻辑
                }
            }
        });
    }

    private void handleWebSocketResponse(ResponseBody responseBody, WebSocketSession session) throws IOException {
        try (BufferedReader reader = new BufferedReader(new InputStreamReader(responseBody.byteStream()))) {
            String line;
            while ((line = reader.readLine()) != null) {
                logger.info("收到SSE数据: {}", line);
                processWebSocketData(line, session);
            }
        } catch (IOException e) {
            logger.error("处理响应数据时出错", e);
        }
    }

    private void processWebSocketData(String data, WebSocketSession session) {
        if (data.startsWith("data: ")) {
            String jsonData = data.substring(6).trim();

            if ("[DONE]".equals(jsonData)) {
                logger.info("数据流已结束");
                return;
            }

            try {
                SseResponse responsePart = JSON.parseObject(jsonData, SseResponse.class);
                session.sendMessage(new TextMessage(JSON.toJSONString(responsePart)));  // 发送数据到客户端
            } catch (Exception e) {
                logger.error("解析或推送 JSON 数据时出错", e);
            }
        }
    }
}
